#include <iostream>
#include <vector>

#include <map>
#include <set>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <matrix.h>
#include <cluster.h>
#include <segment.h>

#include <lemon/smart_graph.h>
#include <lemon/unionfind.h>
#include <lemon/random.h>
#include <lemon/matching.h>

using namespace std;
using namespace sds;
using namespace lemon;

#if !defined(INIT_NORMAL) && !defined(INIT_ADV)
#define INIT_ADV
#endif

#if !defined(JOIN_NORMAL) && !defined(JOIN_MATCHING)
#define JOIN_MATCHING
#endif

int smallClusterSize = 240; //256
double smallDistance = 1.20;

double maxDistance = 106.0;

int stepNum = 10;

double initialCluster = 160;

void Read(const string& filename, const std::string& segname,
	  SDataView<Point>& src, SDataView<int>& seg, 
          vector<Cluster>& sign, vector<SInteger>& size) {

  SFile segfile(segname);
  SDataView<unsigned char> segcol = 
    segfile.getDataset(0).getDataView<unsigned char>();

  size = segcol.getSize();
  
  seg = SDataView<int>(size);

  SDataView<unsigned char> source[CH];
  SFile file(filename);
  for (int k = 0; k < CH; ++k) {
    SDataset dataset = file.getDataset(k);
    source[k] = dataset.getDataView<unsigned char>();
  }

  src = SDataView<Point>(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      src[i][j] = point;
    }
  }  

  ReadSegments(src, segcol, seg, sign);
}

void Write(const string& trgname, const SDataView<int>& seg,
	   std::vector<Cluster> sign, std::vector<Point>& centers, 
	   vector<SInteger>& size) {

  SFile trgfile(trgname, SFile::CREATE);
  SDataset trg = trgfile.createDataset("clustered", UINT8, size);

  std::vector<char> clusters(sign.size());

  for (int i = 0; i < (int)sign.size(); ++i) {
    int minPos = 0;
    double minVal = mon_distance_avg(sign[i], centers[0]);
    for (int k = 1; k < (int)centers.size(); ++k) {
      double val = mon_distance_avg(sign[i], centers[k]);
      if (val < minVal) {
	minVal = val;
	minPos = k;
      }
    }
    if (distance_avg(sign[i], centers[minPos]) < maxDistance) {
      clusters[i] = minPos + 1;
    } else {
      clusters[i] = 0;
    }
  }

  SDataView<unsigned char> clu(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (seg[i][j] != -1) {
	clu[i][j] = clusters[seg[i][j]];
      } else {
	clu[i][j] = 0;
      }
    }
  }
  trg.setDataView<unsigned char>(clu);  
}

#if defined(INIT_NORMAL)
void start(const vector<Cluster>& sign, std::vector<Point>& centers) {
  Cluster info;
  for (int i = 0; i < (int)sign.size(); ++i) {
    info += sign[i];
  }
  Vector expected(CH);
  Matrix covariance(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      covariance(i, j) = info.covariance(i, j);
    }
    expected(i) = info.expected(i);
  }
  Matrix cholesky = covariance.cholesky();
  for (int k = 0; k < initialCluster ++k) {
    Vector v(CH);
    for (int i = 0; i < CH; ++i) {
      v(i) = lemon::rnd.gauss();
    }
    v = v * cholesky + expected;
    Point p;
    for (int i = 0; i < CH; ++i) {
      p[i] = v(i);
    }
    centers.push_back(p);
  }
}

#elif defined(INIT_ADV)

void start(const vector<Cluster>& sign, std::vector<Point>& centers) {

  double dss;
  std::vector<double> ds(sign.size());

  {
    int sn = 0;
    for (int i = 0; i < sign.size(); ++i) {
      sn += sign[i].size(); 
    }

    int r = lemon::rnd[sn];
    int n;
    for (n = 0; ; ++n) {
      if ( r < sign[n].size()) break;
      r -= sign[n].size();
    }

    Point cntr = sign[n].center();
    centers.push_back(cntr);
    
    dss = 0.0;

    for (int i = 0; i < sign.size(); ++i) {
      ds[i] = mon_distance_avg(sign[i], cntr);
      dss += ds[i] * sign[i].size();
    }
  }

  for (int k = 1; k < initialCluster; ++k) {
    std::cerr << k << std::endl;
    double r = lemon::rnd(dss);

    int n;
    for (n = 0; ; ++n) {
      if ( r < ds[n] * sign[n].size()) break;
      r -= ds.data()[n] * sign[n].size();
    }

    Point cntr = sign[n].center();
    centers.push_back(cntr);

    dss = 0.0;
    for (int i = 0; i < sign.size(); ++i) {
      ds[i] = std::min(ds[i], mon_distance_avg(sign[i], cntr));
      dss += ds[i] * sign[i].size();
    }
  }
}

#endif

void classify(const std::vector<Cluster>& sign,
	      std::vector<Point>& centers,
	      std::vector<Cluster>& clusters) {
  clusters.clear();
  clusters.resize(centers.size());
  for (int i = 0; i < (int)sign.size(); ++i) {
    int minPos = 0;
    double minVal = mon_distance_avg(sign[i], centers[0]);
    for (int k = 1; k < (int)centers.size(); ++k) {
      double val = mon_distance_avg(sign[i], centers[k]);
      if (val < minVal) {
	minVal = val;
	minPos = k;
      }
    }
    if (distance_avg(sign[i], centers[minPos]) < maxDistance) {
      clusters[minPos] += sign[i];
    } else {
      clusters.push_back(sign[i]);
      centers.push_back(sign[i].center());
    }
  }
}

class IntMap : public std::vector<int> {
public:
  typedef int Key;

  IntMap(int size, int value) : std::vector<int>(size, value) {}
  void set(int key, int value) { (*this)[key] = value;}
};

void removeSmallClusters(std::vector<Cluster>& infos) {
  for (int i = 0; i < (int)infos.size(); ++i) {
    if (infos[i].size() < smallClusterSize) {
      infos[i] = infos.back();
      infos.pop_back();
      --i;
      std::cout << "Remove: " << infos.size() << std::endl;
    }
  }
}

#if defined(JOIN_NORMAL)

void joinClusters(std::vector<Cluster>& infos, 
		  std::vector<Point>& centers) {
  centers.clear();
  IntMap unionfindMap(infos.size(), -1);
  UnionFind<IntMap > unionfind(unionfindMap);

  for (int i = 0; i < (int)infos.size(); ++i) {
    unionfind.insert(i);
  }
  
  for (int i = 0; i < (int)infos.size(); ++i) {
    for (int j = i + 1; j < (int)infos.size(); ++j) {
      if (distance_bat(infos[i], infos[j]) < smallDistance) {
	unionfind.join(i, j);
	std::cout << "Join: " << distance_bat(infos[i], infos[j]) << std::endl;
      }
    }
  }
  std::map<int, Cluster> nc;
  for (int i = 0; i < (int)infos.size(); ++i) {
    nc[unionfind.find(i)] += infos[i];
  }
  std::map<int, Cluster>::iterator it;
  for (it = nc.begin(); it != nc.end(); ++it) {
    centers.push_back(it->second.center());
  }
}

#elif defined(JOIN_MATCHING)

void joinClusters(std::vector<Cluster>& infos, 
		  std::vector<Point>& centers) {

  lemon::SmartGraph graph;
  std::vector<lemon::SmartGraph::Node> nodes(infos.size());
  lemon::SmartGraph::EdgeMap<double> weight(graph);

  centers.clear();

  for (int i = 0; i < (int)infos.size(); ++i) {
    nodes[i] = graph.addNode();
  }

  for (int i = 0; i < (int)infos.size(); ++i) {
    for (int j = i + 1; j < (int)infos.size(); ++j) {
      double dist = distance_bat(infos[i], infos[j]);
      if (dist < smallDistance) {
	lemon::SmartGraph::Edge ue = graph.addEdge(nodes[i], nodes[j]);
	weight[ue] = smallDistance - dist;
      }
    }
  }
  lemon::MaxWeightedMatching<lemon::SmartGraph,
    lemon::SmartGraph::EdgeMap<double> > mwm(graph, weight);
  mwm.run();

  for (int i = 0; i < (int)infos.size(); ++i) {
    if (mwm.matching(nodes[i]) != INVALID) {
      int mi = graph.id(mwm.mate(nodes[i])); 
      if (mi < i) continue;
      Cluster merge = infos[i] + infos[mi];
      centers.push_back(merge.center());
    } else {
      centers.push_back(infos[i].center());
    }
  }
  std::cout << std::endl;
}


#endif


int main(int argc, const char *argv[]) {
  assert(argc == 4);

  SDataView<int> seg;
  SDataView<Point> src;
  vector<SInteger> size;

  vector<Cluster> sign;

  Read(argv[1], argv[2], src, seg, sign, size);
  std::cerr << "Reading done." << std::endl; 

  std::vector<Point> centers;

  start(sign, centers);
  
  for (int step = 0; step < stepNum; ++step) {
    std::cout << (double)step / stepNum << std::endl;
    std::vector<Cluster> clusters;
    classify(sign, centers, clusters);
    if (step > 0) {
      removeSmallClusters(clusters);
      joinClusters(clusters, centers);
    } else {
      for (int i = 0; i < clusters.size(); ++i) {
        centers[i] = clusters[i].center();
      } 
    }
    std::cout << "Num: " << centers.size() << std::endl;
  }
  
  Write(argv[3], seg, sign, centers, size);
  return 0;
}
